# _*_ coding: utf-8 _*_
"""
Time:     2024/9/14
Author:   2051261
File:     base_dest_entity.py
"""
from __future__ import annotations
import typing as t

from schematics.models import Model
from schematics.types import StringType
from schematics.types import ListType
from schematics.types import BooleanType
from schematics.types import IntType
from schematics.types import DictType
from schematics.types import BaseType

from extractor.const import ContextKey
from extractor.const import ConstKey
from extractor.const import Mark
from mine.db.db_operate import db_operate
from extractor.model.factory import SQLBuilderFactory
from extractor.flow.base_node_flow import BaseNodeFlow

if t.TYPE_CHECKING:
    from extractor.flow.node_model.log_node import RecordLog
    from extractor.flow.task_entry import Cache


class BaseDestEntityModel(Model):
    conn_name: str = StringType(required=True)
    table_name: str = StringType(required=True)

    dirty_conn_name: str = StringType()
    dirty_unique_key: list[str] = ListType(StringType())
    dirty_update_key: list[str] = ListType(StringType())

    insert_version: bool = IntType()  # 1 表示同步插入version 2 表示 只有主表发生插入或者更新才插入version

    ref_column: str = StringType()
    ref_from_column: str = StringType()
    pop_columns: list[str] = ListType(StringType())  # 比如清洗库 znz_id 和 znz_create 不继承源数据,需要配置移除,

    unique_key: list[str] = ListType(StringType())
    update_key: list[str] = ListType(StringType())
    extra_update: dict = DictType(BaseType)
    replace: bool = BooleanType(default=False)
    ignore: bool = BooleanType(default=False)
    znz_mark: str = StringType(default='znz_mark')


def _insert_dirty_db(dirty: dict, model: BaseDestEntityModel):
    """

    :param dirty:
    :param model:
    :return:
    """
    dirty_conn_name, table_name = model.dirty_conn_name, model.table_name

    dirty_unique, dirty_update = model.dirty_unique_key, model.dirty_update_key
    db_type = db_operate.get_db_type(dirty_conn_name)
    sql_builder = SQLBuilderFactory().get_query_builder(db_type)
    sql = sql_builder.insert_sql(table_name, dirty, unique=dirty_unique, update=dirty_update, replace=True)

    return db_operate.insert_one(dirty_conn_name, sql, dirty)


def _insert_ver_db(data: dict, model: BaseDestEntityModel):
    """

    :param data:
    :param model:
    :return:
    """
    conn_name, table_name = model.conn_name, model.table_name
    table_name = table_name + "_ver"
    unique = model.unique_key
    update = model.update_key
    db_type = db_operate.get_db_type(conn_name)
    sql_builder = SQLBuilderFactory().get_query_builder(db_type)
    sql = sql_builder.insert_sql(table_name, data, unique=unique, update=update)

    return db_operate.insert_one(conn_name, sql, data)


def _insert_dest_db(data, model: BaseDestEntityModel):
    """

    :param data:
    :param model:
    :return:
    """
    conn_name, table_name = model.conn_name, model.table_name
    unique, update, extra_update = model.unique_key, model.update_key, model.extra_update
    combined_data = {**data}  # 创建 data 的浅拷贝
    if extra_update:
        combined_data.update(extra_update)
        update = update + list(extra_update.keys()) if isinstance(update, list) else update
    db_type = db_operate.get_db_type(conn_name)
    replace, ignore = model.replace, model.ignore
    sql_builder = SQLBuilderFactory().get_query_builder(db_type)
    sql = sql_builder.insert_sql(table_name, combined_data,
                                 unique=unique, update=update,
                                 replace=replace, ignore=ignore)
    return db_operate.insert_one(conn_name, sql, combined_data)


def _get_table_all_columns(conn_name: str, table_name: str, cache: Cache):
    db_type = db_operate.get_db_type(conn_name)
    sql_builder = SQLBuilderFactory().get_query_builder(db_type)
    sql = sql_builder.get_table_columns(table_name)
    result = db_operate.execute_query(conn_name, sql)
    columns = [column[0].lower() for column in result]
    cache.cache_table_columns(table_name, columns)


def _get_data_by_columns(data: dict, columns: list[str]) -> dict:
    lower_data = {k.lower(): v for k, v in data.items()}

    # 遍历 columns 列表，尝试匹配列（不区分大小写）
    result = {}
    for col in columns:
        if col in data:  # 如果列名完全匹配，直接取值
            result[col] = data[col]
        elif col.lower() in lower_data:  # 如果小写列名匹配，取小写匹配的值
            result[col] = lower_data[col.lower()]

    return result


def _pop_dict_key(data: dict, columns: list[str]):
    for col in columns:
        data.pop(col.lower())


def _insert_db_version_methods(data: dict, result: dict, model: BaseDestEntityModel, cache: Cache):
    conn_name, table_name = model.conn_name, model.table_name
    if cache.get_table_columns() is None or cache.get_table_columns().get(table_name + "_ver") is None:
        _get_table_all_columns(conn_name, table_name + "_ver", cache)
    columns = cache.get_table_columns()[table_name + "_ver"]
    insert_data = _get_data_by_columns(data, columns)
    if 1 == model.insert_version:
        _insert_ver_db(insert_data, model)
    elif 2 == model.insert_version and result[ConstKey.LAST_ROW_ID] != 0:
        _insert_ver_db(insert_data, model)


def _handle_insert_table_result(result: t.Union[dict, int], record: t.DefaultDict,
                                ref_id_value: int, model: BaseDestEntityModel):
    if result in [Mark.UNIQUE_CONFLICT.value, Mark.INSERT_FAILED.value]:
        if result == Mark.UNIQUE_CONFLICT.value:
            record["conflict"] += 1
        else:
            record["failed"] += 1
        if model.dirty_conn_name is not None:
            dirty = {model.znz_mark: result, model.ref_from_column: ref_id_value}
            _insert_dirty_db(dirty, model)
    elif result[ConstKey.AFFECTED_ROWS] == 1 and result[ConstKey.LAST_ROW_ID]:
        record["insert"] += 1
    elif result[ConstKey.AFFECTED_ROWS]:
        record["update"] += 1


class BaseDestEntityNode(BaseNodeFlow):

    async def __call__(self, record_log: RecordLog, context: t.Iterator[t.Dict[str, t.Any]],
                 cache: Cache, record: t.DefaultDict, kwargs=None) -> None:
        """

        :param record_log: 日志记录
        :param context: 上下文
        :param cache: 缓存
        :param record: 操作记录
        :param kwargs: 函数必要参数 参数列表参考 BaseDestEntityModel 模型
        :return:
        """
        model = BaseDestEntityModel(kwargs)
        conn_name, table_name, insert_version = model.conn_name, model.table_name, model.insert_version
        if await cache.get_table_columns() is None or await cache.get_table_columns().get(table_name) is None:
            _get_table_all_columns(conn_name, table_name, cache)
        table_columns = cache.get_table_columns()[table_name]
        for item in context:
            dirty, data = item.get(ContextKey.DIRTY), item[ContextKey.DATA]
            ref_id_value = data.get(model.ref_from_column)
            value = ref_id_value if await cache.get_key() is not None else None
            record["total"] += 1
            if dirty:
                # 有脏数据并且配置了赃库
                if model.dirty_conn_name is not None:
                    _insert_dirty_db(dirty, model)
                    record["dirty"] += 1
                else:
                    record["skip"] += 1
            else:
                ref_col, ref_from_col = model.ref_column, model.ref_from_column
                if ref_col is not None and ref_from_col is not None:
                    data[ref_col] = data[ref_from_col]
                if model.pop_columns:
                    _pop_dict_key(data, model.pop_columns)
                insert_data = _get_data_by_columns(data, table_columns)
                result = _insert_dest_db(insert_data, model)
                _handle_insert_table_result(result, record, ref_id_value, model)
                # 是否需要插入版本库
                if model.insert_version and result != Mark.INSERT_FAILED.value:
                    _insert_db_version_methods(data, result, model, cache)
            # 更新缓存
            if value is not None:
                await cache.update(value)
